package mr

import "fmt"
import "log"
import "time"
import "os"
import "strings"
import "io/ioutil"
import "net/rpc"
import "hash/fnv"
import "sort"
import "strconv"

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }
//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}


//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.

	// uncomment to send the Example RPC to the coordinator.
	// CallExample()
	for ; ; {
		task := AssignedTask{}
		rep := ExampleReply{}
		ok := call("Coordinator.AssignTask", &rep, &task)
//		log.Printf("node received assigned task %v\n", &task)
		if (!ok) {
			return
		}
		switch(task.TaskType) {
		case MAP_TASK:
			HandleMapTask(mapf, reducef, &task)
			continue
		case REDUCE_TASK:
			HandleReduceTask(mapf, reducef, &task)
			continue
		case TEMP_NO_TASK:
			time.Sleep(1 * time.Second)
			continue
		default:
			return
		}
	}
}

func HandleMapTask(mapf func(string, string) []KeyValue,
reducef func(string, []string) string, task *AssignedTask) {
	filename := task.InputFileNames[0]
	task_id := task.TaskId
	reduce_num := task.NumOfOutput

	// Read file, this code is borrowed from mrsequential
	file, err := os.Open(filename)
	if err != nil {
		log.Fatalf("cannot open %v", filename)
	}
	content, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatalf("cannot read %v", filename)
	}
	file.Close()
	kva := mapf(filename, string(content))

	file_arr := make([]*os.File, reduce_num)

	for i := 0; i < reduce_num; i++ {
		temp_file, err := ioutil.TempFile("", "temp-*.txt")
		if err != nil {
			log.Fatalf("cannot create temp file")
		}
		file_arr[i] = temp_file
	}

	for _, pair := range kva {
		idx := ihash(pair.Key) % reduce_num
		fmt.Fprintf(file_arr[idx], "%v %v\n", pair.Key, pair.Value)
	}

	finish_task := FinishedTask{}
	finish_task.TaskType = MAP_TASK
	finish_task.TaskId = task_id
	finish_task.OutputFileNames = make([]string, reduce_num)

	for i := 0; i < reduce_num; i++ {
		immedia_file_name := "mr-" + strconv.Itoa(task_id) + "-" + strconv.Itoa(i)
		err := os.Rename(file_arr[i].Name(), immedia_file_name)
		if err != nil {
			log.Fatalf("cannot rename %v to %v", file_arr[i].Name(), immedia_file_name)
		}
		finish_task.OutputFileNames[i] = immedia_file_name
	}

	rep := ExampleReply{}

	ok := call("Coordinator.TaskCompleted", &finish_task, &rep)
	if !ok {
		log.Fatalf("submit task error...")
	}
}

func HandleReduceTask(mapf func(string, string) []KeyValue,
reducef func(string, []string) string, task *AssignedTask) {
	task_id := task.TaskId
	intermediate := []KeyValue{}
	for _, filename := range task.InputFileNames {
		file, err := os.Open(filename)
		if err != nil {
			log.Fatalf("cannot open %v", filename)
		}
		content, err := ioutil.ReadAll(file)
		if err != nil {
			log.Fatalf("cannot read %v", filename)
		}
		file.Close()
		str_content := string(content)
		r:= strings.NewReader(str_content)
		for ; ; {
			pair := KeyValue{}
			n, err := fmt.Fscanf(r, "%v %v\n", &pair.Key, &pair.Value)
			if (n == 0 || err != nil) {
				break
			}
			intermediate = append(intermediate, pair)
		}
	}

	sort.Sort(ByKey(intermediate))

	oname := "mr-out-" + strconv.Itoa(task_id)
	ofile, _ := os.Create(oname)

	i := 0
	for i < len(intermediate) {
		j := i + 1
		for j < len(intermediate) && intermediate[j].Key == intermediate[i].Key {
			j++
		}
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, intermediate[k].Value)
		}
		output := reducef(intermediate[i].Key, values)

		// this is the correct format for each line of Reduce output.
		fmt.Fprintf(ofile, "%v %v\n", intermediate[i].Key, output)

		i = j
	}

	ofile.Close()

	finish_task := FinishedTask{}
	finish_task.TaskType = REDUCE_TASK
	finish_task.TaskId = task_id

	rep := ExampleReply{}

	ok := call("Coordinator.TaskCompleted", &finish_task, &rep)
	if !ok {
		log.Fatalf("submit task error...")
	}
}
//
// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	// the "Coordinator.Example" tells the
	// receiving server that we'd like to call
	// the Example() method of struct Coordinator.
	ok := call("Coordinator.Example", &args, &reply)
	if ok {
		// reply.Y should be 100.
		fmt.Printf("reply.Y %v\n", reply.Y)
	} else {
		fmt.Printf("call failed!\n")
	}
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
